import base64
import hashlib
import json
import os
import socket
import threading
import time
import tornado.ioloop
import tornado.web
import tornado.websocket
import tornado.gen
from tornado.httpclient import AsyncHTTPClient, HTTPClientError
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.backends import default_backend
import rsa
import base64  
class P2PTornadoCommunication:
    def __init__(self):
        self.backend = default_backend()
        self.session_key = None
        self.session_iv = None
        self.application = None
        self.http_server = None
        self.ioloop = None
        self.server_thread = None
        self.peers = {}
        self.messages = []
        self.message_callback = None
        self.is_server_running = False
        
        # 确保有RSA密钥对
        self.ensure_rsa_keys()
        
    def ensure_rsa_keys(self):
        """确保存在RSA密钥对"""
        self.public_key_path = os.path.join(os.getcwd(), "public_key.pem")
        self.private_key_path = os.path.join(os.getcwd(), "private_key.pem")
        
        if not os.path.exists(self.public_key_path) or not os.path.exists(self.private_key_path):
            # 生成新的RSA密钥对
            print("生成新的RSA密钥对...")
            public_key, private_key = rsa.newkeys(2048)
            
            with open(self.public_key_path, 'wb') as f:
                f.write(public_key.save_pkcs1())
            with open(self.private_key_path, 'wb') as f:
                f.write(private_key.save_pkcs1())
        
        # 加载RSA密钥
        with open(self.public_key_path, 'rb') as f:
            self.public_key = rsa.PublicKey.load_pkcs1(f.read())
        with open(self.private_key_path, 'rb') as f:
            self.private_key = rsa.PrivateKey.load_pkcs1(f.read())
    
    def start_server(self, port=8080):
        """启动Tornado服务器"""
        self.port = port
        self.server_address = ('', port)
        
        class MainHandler(tornado.web.RequestHandler):
            def initialize(self, communicator):
                self.communicator = communicator
                
            def get(self):
                self.write("点对点加密通信服务器正在运行")
        
        class StatusHandler(tornado.web.RequestHandler):
            def initialize(self, communicator):
                self.communicator = communicator
                
            def get(self):
                self.set_header("Content-Type", "application/json")
                self.write(json.dumps({"status": "online"}))
        
        class PublicKeyHandler(tornado.web.RequestHandler):
            def initialize(self, communicator):
                self.communicator = communicator
                
            def get(self):
                self.set_header("Content-Type", "application/octet-stream")
                with open(self.communicator.public_key_path, 'rb') as f:
                    self.write(f.read())
        
        class MessageHandler(tornado.web.RequestHandler):
            def initialize(self, communicator):
                self.communicator = communicator
                
            async def post(self):
                try:
                    # 获取请求体
                    encrypted_data = base64.b64decode(self.request.body)
                    # self.communicator.session_key:是当前会话的密钥
                    print(f"会话密钥: {self.communicator.session_key}")
                    if self.communicator.session_key:
                        print("使用会话密钥解密")
                        # 使用会话密钥解密
                        decrypted_message = self.communicator._decrypt_with_session_key(encrypted_data)
                        print(f"解密后的消息: {decrypted_message}")
                        json_str = decrypted_message.replace("'", "\"")
                        # 解析消息str to dict
                        decrypted_message = json.loads(json_str)
                        print(f"解析后的断开连接请求: {decrypted_message}")
                        message = decrypted_message['message']
                        if message == '断开连接':
                            peer_ip = decrypted_message['peer_ip']
                            peer_port = decrypted_message['peer_port']
                            print(f"收到断开连接请求: {peer_ip}:{peer_port}")
                        
                            # 断开连接：需要删除 peers 中的记录
                            # 删除session_key
                        
                            self.communicator.session_key = None
                            self.communicator.session_iv = None
                            self.communicator.peers.pop(f"{peer_ip}:{peer_port}", None)
                            # 关闭所有WebSocket连接
                            for conn in CommunicationWebSocket.connections:
                                conn.close()
                            print(f"成功断开连接到对等节点: {peer_ip}:{peer_port}")
                        
                    else:
                        print("初始消息，使用RSA解密获取会话密钥")
                        # 初始消息，使用RSA解密获取会话密钥
                        decrypted_data = rsa.decrypt(encrypted_data, self.communicator.private_key)
                        session_info = json.loads(decrypted_data.decode('utf-8'))
                        self.communicator.session_key = base64.b64decode(session_info['session_key'])
                        self.communicator.session_iv = base64.b64decode(session_info['session_iv'])
                        decrypted_message = session_info['message']
                        print(f"解密后的消息: {decrypted_message}")
                    
                    # 保存消息
                    peer_ip = self.request.remote_ip
                    message_data = {
                        'from': peer_ip,
                        'message': decrypted_message,
                        'timestamp': time.time()
                    }
                    self.communicator.messages.append(message_data)
                    
                    # 调用回调函数
                    if self.communicator.message_callback:
                        self.communicator.message_callback(message_data)
                    
                    
                    self.set_header("Content-Type", "application/json")
                    self.write(json.dumps({"status": "success"}))
                except Exception as e:
                    self.set_status(500)
                    self.set_header("Content-Type", "application/json")
                    self.write(json.dumps({"status": "error", "message": str(e)}))
        
        # 创建WebSocket处理程序
        class CommunicationWebSocket(tornado.websocket.WebSocketHandler):
            connections = set()
            
            def initialize(self, communicator):
                self.communicator = communicator
                self.peer_ip = None
                
            def open(self):
                self.peer_ip = self.request.remote_ip
                CommunicationWebSocket.connections.add(self)
                print(f"WebSocket连接已建立: {self.peer_ip}")
                
            def on_message(self, message):
                try:
                    # 解密消息
                    encrypted_data = base64.b64decode(message)
                    decrypted_message = self.communicator._decrypt_with_session_key(encrypted_data)
                    
                    # 保存消息
                    message_data = {
                        'from': self.peer_ip,
                        'message': decrypted_message,
                        'timestamp': time.time()
                    }
                    self.communicator.messages.append(message_data)
                    
                    # 调用回调函数
                    if self.communicator.message_callback:
                        self.communicator.message_callback(message_data)
                except Exception as e:
                    print(f"处理WebSocket消息时出错: {str(e)}")
                    self.write_message(base64.b64encode(str(e).encode('utf-8')))
                
            def on_close(self):
                CommunicationWebSocket.connections.remove(self)
                print(f"WebSocket连接已关闭: {self.peer_ip}")
                
            def check_origin(self, origin):
                # 允许所有来源的WebSocket连接
                return True
        
        # 创建应用程序
        self.application = tornado.web.Application([
            (r'/', MainHandler, dict(communicator=self)),
            (r'/status', StatusHandler, dict(communicator=self)),
            (r'/public_key', PublicKeyHandler, dict(communicator=self)),
            (r'/message', MessageHandler, dict(communicator=self)),
            (r'/ws', CommunicationWebSocket, dict(communicator=self)),
        ])
        
        # 创建HTTP服务器
        self.http_server = self.application.listen(port)
        self.ioloop = tornado.ioloop.IOLoop.current()
        
        # 在单独线程中运行服务器
        self.server_thread = threading.Thread(target=self._run_server)
        self.server_thread.daemon = True
        self.server_thread.start()
        self.is_server_running = True
        
        print(f"Tornado服务器已启动，监听端口 {port}")
    
    def _run_server(self):
        """在单独线程中运行服务器"""
        try:
            self.ioloop.start()
        except KeyboardInterrupt:
            pass
    
    def stop_server(self):
        """停止Tornado服务器"""
        if self.ioloop and self.is_server_running:
            self.ioloop.add_callback(lambda x: x.stop(), self.ioloop)
            self.server_thread.join(timeout=1.0)
            self.is_server_running = False
            print("Tornado服务器已停止")
    
    async def connect_to_peer_async(self, peer_ip, peer_port=8080):
        """异步连接到对等节点并交换密钥"""
        peer_url = f"http://{peer_ip}:{peer_port}"
        print(f"尝试连接到对等节点: {peer_ip}:{peer_port}",peer_url)
        try:
            # 检查对等节点是否在线
            client = AsyncHTTPClient()
            print("client:",client)
            # 检查对等节点是否在线
            response = await client.fetch(f"{peer_url}/status")
            print(f"获取对等节点状态响应: {response.code}")
            if response.code != 200:
                raise Exception("对等节点不在线")
            
            # 获取对等节点的公钥
            response = await client.fetch(f"{peer_url}/public_key")
            print(f"获取对等节点公钥响应: {response.code}")
            peer_public_key = rsa.PublicKey.load_pkcs1(response.body)
            print(f"获取对等节点公钥: {peer_public_key}")
            # 生成会话密钥和IV
            if not self.session_key:
                self.session_key = os.urandom(32)  # 256位AES密钥
                self.session_iv = os.urandom(16)  # 16字节IV
            print(f"会话密钥: {self.session_key}")
            # 将会话密钥和初始消息加密并发送给对等节点
            session_info = {
                'session_key': base64.b64encode(self.session_key).decode('utf-8'),
                'session_iv': base64.b64encode(self.session_iv).decode('utf-8'),
                'message': "建立安全连接成功！"
            }
            print(f"会话信息: {session_info}")
            
            # 使用对等节点的公钥加密会话信息
            encrypted_data = rsa.encrypt(
                json.dumps(session_info).encode('utf-8'),
                peer_public_key
            )
            print(f"加密数据: {encrypted_data}")
            # 发送加密数据
            response = await client.fetch(
                f"{peer_url}/message",
                method='POST',
                body=base64.b64encode(encrypted_data),
                headers={'Content-Type': 'application/octet-stream'},
                # timeout=5
            )
            print(f"发送加密数据响应: {response.code}")
            if response.code == 200:
                self.peers[f"{peer_ip}:{peer_port}"] = {
                    'ip': peer_ip,
                    'port': peer_port,
                    'public_key': peer_public_key,
                    'connected_at': time.time()
                }
                print(f"成功连接到对等节点: {peer_ip}:{peer_port}")
                return True
            else:
                raise Exception(f"连接失败: {response.body.decode('utf-8')}")
        except HTTPClientError as e:
            raise Exception(f"HTTP错误: {str(e)}")
        except Exception as e:
            print(f"连接到对等节点失败: {str(e)}")
            return False
    
    def connect_to_peer(self, peer_ip, peer_port=8080):
        """连接到对等节点并交换密钥（同步接口）"""
        # 创建一个新的IOLoop用于同步操作
        io_loop = tornado.ioloop.IOLoop()
        result = [False]
        
        async def connect_task():
            result[0] = await self.connect_to_peer_async(peer_ip, peer_port)
            
        def run_sync():
            io_loop.run_sync(connect_task)
        
        try:
            # 在单独的线程中运行异步连接任务
            thread = threading.Thread(target=run_sync)
            thread.start()
            thread.join(timeout=10.0)  # 设置超时
            
            if thread.is_alive():
                print("连接超时")
                return False
            
            return result[0]
        finally:
            # 清理资源
            io_loop.close()
    
    async def disconnect_peer_async(self, peer_ip, peer_port):
        """异步断开与对等节点的连接"""
        peer_url = f"http://{peer_ip}:{peer_port}"
        print(f"尝试断开连接到对等节点: {peer_ip}:{peer_port}",peer_url)
        message = {
            "peer_ip":peer_ip,
            "peer_port":peer_port,
            "message": "断开连接"
        }.__str__()
        try:
          # 使用会话密钥加密消息
            encrypted_message = self._encrypt_with_session_key(message)
            
            # 发送加密消息
            client = AsyncHTTPClient()
            response = await client.fetch(
                f"{peer_url}/message",
                method='POST',
                body=base64.b64encode(encrypted_message),
                headers={'Content-Type': 'application/octet-stream'},
                # timeout=5
            )
            
            if response.code == 200:
                self.peers.pop(f"{peer_ip}:{peer_port}", None)
                self.session_key = None
                self.session_iv = None
                print(f"成功断开连接到对等节点: {peer_ip}:{peer_port}")
                # 保存发送的消息
                message_data = {
                    'to': f"{peer_ip}:{peer_port}",
                    'message': message,
                    'timestamp': time.time()
                }
                self.messages.append(message_data)
                return True
            else:
                raise Exception(f"断开连接失败: {response.body.decode('utf-8')}")
        except HTTPClientError as e:
            raise Exception(f"HTTP错误: {str(e)}")
        except Exception as e:
            print(f"断开连接到对等节点失败: {str(e)}")
            return False
         
    def disconnect_peer(self, peer_ip, peer_port):
        """断开与对等节点的连接"""
        # 创建一个新的IOLoop用于同步操作
        io_loop = tornado.ioloop.IOLoop()
        result = [False]
        
        async def disconnect_task():
            result[0] = await self.disconnect_peer_async(peer_ip, peer_port)
            
        def run_sync():
            io_loop.run_sync(disconnect_task)
        
        try:
            # 在单独的线程中运行异步连接任务
            thread = threading.Thread(target=run_sync)
            thread.start()
            thread.join(timeout=10.0)  # 设置超时
            
            if thread.is_alive():
                print("连接超时")
                return False
            
            return result[0]
        finally:
            # 清理资源
            io_loop.close()

    async def send_message_async(self, peer_ip, peer_port, message):
        """异步发送加密消息给对等节点"""
        if not self.session_key:
            raise Exception("尚未建立安全连接")
        
        peer_url = f"http://{peer_ip}:{peer_port}"
        
        try:
            message = {
                "message": message
            }.__str__()
            # 使用会话密钥加密消息
            encrypted_message = self._encrypt_with_session_key(message)
            
            # 发送加密消息
            client = AsyncHTTPClient()
            response = await client.fetch(
                f"{peer_url}/message",
                method='POST',
                body=base64.b64encode(encrypted_message),
                headers={'Content-Type': 'application/octet-stream'},
                # timeout=5
            )
            
            if response.code == 200:
                # 保存发送的消息
                message_data = {
                    'to': f"{peer_ip}:{peer_port}",
                    'message': message,
                    'timestamp': time.time()
                }
                self.messages.append(message_data)
                return True
            else:
                raise Exception(f"发送消息失败: {response.body.decode('utf-8')}")
        except Exception as e:
            print(f"发送消息失败: {str(e)}")
            return False
    
    def send_message(self, peer_ip, peer_port, message):
        """发送加密消息给对等节点（同步接口）"""
        # 创建一个新的IOLoop用于同步操作
        io_loop = tornado.ioloop.IOLoop()
        result = [False]
        
        async def send_task():
            result[0] = await self.send_message_async(peer_ip, peer_port, message)
        
        def run_sync():
            io_loop.run_sync(send_task)
        
        try:
            # 在单独的线程中运行异步发送任务
            thread = threading.Thread(target=run_sync)
            thread.start()
            thread.join(timeout=10.0)  # 设置超时
            
            if thread.is_alive():
                print("发送消息超时")
                return False
            
            return result[0]
        finally:
            # 清理资源
            io_loop.close()
    
    def _encrypt_with_session_key(self, data):
        """使用会话密钥加密数据"""
        if not self.session_key or not self.session_iv:
            raise Exception("会话密钥未初始化")
        
        # 确保数据是字节
        if isinstance(data, str):
            data = data.encode('utf-8')
        
        # 填充数据
        padder = padding.PKCS7(algorithms.AES.block_size).padder()
        padded_data = padder.update(data) + padder.finalize()
        
        # 加密数据
        cipher = Cipher(algorithms.AES(self.session_key), modes.CBC(self.session_iv), backend=self.backend)
        encryptor = cipher.encryptor()
        encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
        
        return encrypted_data
    
    def _decrypt_with_session_key(self, encrypted_data):
        """使用会话密钥解密数据"""
        if not self.session_key or not self.session_iv:
            raise Exception("会话密钥未初始化")
        # print(f"会话密钥: {self.session_key}")
        # print(f"会话IV: {self.session_iv}")
        # print(f"加密数据: {encrypted_data}")
        # 解密数据
        cipher = Cipher(algorithms.AES(self.session_key), modes.CBC(self.session_iv), backend=self.backend)
        decryptor = cipher.decryptor()
        padded_data = decryptor.update(encrypted_data) + decryptor.finalize()
        # print(f"解密前的数据: {padded_data}")
        # 去除填充
        unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder()
        data = unpadder.update(padded_data) + unpadder.finalize()
        # 尝试转换为字符串
        try:
            return data.decode('utf-8')
        except:
            return data
    
    def get_local_ip(self):
        """获取本地IP地址"""
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            # 不实际连接，只是为了获取本地IP
            s.connect(('8.8.8.8', 80))
            local_ip = s.getsockname()[0]
        finally:
            s.close()
        return local_ip

# 创建全局实例
p2p_communication = P2PTornadoCommunication()